#!/usr/bin/python3 -u
# Synthesize an OS update by modifying ELF files in a "benign" way
# (adding an ELF note).  This way the upgrade is effectively a no-op,
# but we still test most of the actual mechanics of an upgrade
# such as writing new files, etc.
#
# This uses the latest build's OSTree commit as source, and will
# update the ref but not generate a new coreos-assembler build.

import argparse
import gi
import os
import random
import subprocess
import stat
import sys
import time
import tempfile

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from cosalib.builds import Builds
from cosalib.meta import GenericBuildMeta as Meta

gi.require_version('OSTree', '1.0')
from gi.repository import GLib, Gio, OSTree

# There are ELF files outside of these paths, but we don't
# care
SUBDIRS = ["/usr/" + x for x in ["bin", "sbin", "lib", "lib/systemd", "lib64"]]

parser = argparse.ArgumentParser()
parser.add_argument("--repo", help="OSTree repo path", default='tmp/repo')
parser.add_argument("--src-ref", help="Branch to use as source for update")
parser.add_argument("--ref", help="Branch to target for update (default is build ref)")
parser.add_argument("--initramfs", help="Generate an update for the initramfs", default=True)
parser.add_argument("--percentage", help="Approximate percentage of files to update", default=20, type=int)
args = parser.parse_args()

if args.src_ref is None and args.ref is None:
    build = Meta(build=Builds().get_latest())
    args.src_ref = build['ostree-commit']
    args.ref = build['ref']
if args.src_ref is None:
    args.src_ref = args.ref

version = "synthetic-osupdate-{}".format(int(time.time()))

repo = OSTree.Repo.new(Gio.File.new_for_path(args.repo))
repo.open(None)

[_, root, rev] = repo.read_commit(args.src_ref, None)


def generate_modified_elf_files(srcd, destd, notepath):
    e = srcd.enumerate_children("standard::name,standard::type,unix::mode", Gio.FileQueryInfoFlags.NOFOLLOW_SYMLINKS, None)
    candidates = []
    while True:
        fi = e.next_file(None)
        if fi is None:
            break
        # Must be a regular file greater than 4k
        # in size, readable and executable but not suid
        # and owned by 0:0
        if fi.get_file_type() != Gio.FileType.REGULAR:
            continue
        if fi.get_size() < 4096:
            continue
        if (fi.get_attribute_uint32("unix::uid") != 0 or
           fi.get_attribute_uint32("unix::gid") != 0):
            continue
        mode = fi.get_attribute_uint32("unix::mode")
        if mode & (stat.S_ISUID | stat.S_ISGID) > 0:
            continue
        if not (mode & stat.S_IRUSR > 0):
            continue
        if not (mode & stat.S_IXOTH > 0):
            continue
        candidates.append(fi)
    n_candidates = len(candidates)
    n = (n_candidates * args.percentage) // 100
    targets = 0
    modified_bytes = 0
    while len(candidates) > 0:
        if targets >= n:
            break
        i = random.randrange(len(candidates))
        candidate = candidates[i]
        f = Gio.BufferedInputStream.new(e.get_child(candidate).read(None))
        f.fill(1024, None)
        buf = f.peek_buffer()
        assert len(buf) > 5
        del candidates[i]
        if not (buf[0] == 0x7F and buf[1:4] == b'ELF'):
            continue
        name = candidate.get_name()
        destpath = destd + '/' + name
        outf = Gio.File.new_for_path(destpath).create(0, None)
        outf.splice(f, 0, None)
        outf.close(None)
        try:
            subprocess.check_call(['objcopy', f"--add-section=.note.coreos-synthetic={notepath}", destpath])
        except subprocess.CalledProcessError as e:
            raise Exception(f"Failed to process {destpath}") from e
        os.chmod(destpath, candidate.get_attribute_uint32("unix::mode"))
        modified_bytes += os.stat(destpath).st_size
        targets += 1
    return (targets, n_candidates, modified_bytes)


with tempfile.TemporaryDirectory(prefix='cosa-dev-synth-update') as tmpd:
    # Create a subdirectory so we can use --consume without deleting the
    # parent, which would potentially confuse tempfile
    subd = tmpd + '/c'
    notepath = tmpd + 'note'
    with open(notepath, 'w') as f:
        f.write("Generated by coreos-assembler dev-synthesize-osupdate\n")
    os.makedirs(subd)
    for d in SUBDIRS:
        destd = subd + d
        os.makedirs(destd)
        (m, n, sz) = generate_modified_elf_files(root.get_child(d), destd, notepath)
        print("{}: Modified {}/{} files, {}".format(d, m, n, GLib.format_size(sz)))

    subprocess.check_call(['ostree', f'--repo={args.repo}', 'commit', '--consume',
                           '-b', args.ref, f'--base={args.src_ref}',
                           f'--add-metadata-string=version={version}',
                           f'--tree=dir={subd}', '--owner-uid=0', '--owner-gid=0',
                           '--selinux-policy-from-base', '--table-output',
                           '--link-checkout-speedup', '--no-bindings', '--no-xattrs'])
    print(f"Updated {args.ref}")
